import tensorflow as tf
from tensorflow.keras import layers, losses, optimizers, metrics, activations, models
import numpy as np
from sklearn.model_selection import train_test_split
import os
from python_ai.common.xcommon import *
import matplotlib.pyplot as plt

np.random.seed(777)
tf.random.set_seed(777)

(x_train, y_train), (x_test, y_test) = tf.keras.datasets.cifar10.load_data()
print(x_train.shape, y_train.shape, x_test.shape, y_test.shape)

batch_size = 32
m_train, _, _, _ = x_train.shape
ds = tf.data.Dataset.from_tensor_slices((x_train, y_train))\
    .shuffle(buffer_size=m_train)\
    .batch(batch_size=batch_size)\
    .prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

generator = ds.as_numpy_iterator()

plt.figure(figsize=[12, 6])
spr = 4
spc = 8
spn = 0
for bx, by in generator:
    for i, bxi in enumerate(bx):
        spn += 1
        plt.subplot(spr, spc, spn)
        plt.imshow(bxi)
        plt.axis('off')
        plt.title(str(by[i, 0]))
    break
